#!/usr/bin/env python3
###
# @Author: Logan.Li
# @Gitee: https://gitee.com/attacker
# @email: admin@attacker.club
# @Date: 2025-03-14 08:25:30
# @LastEditTime: 2025-03-14 08:25:35
# @Description: AWS CloudWatch监控告警脚本
###

import os
import sys
import json
import logging
import configparser
from datetime import datetime, timedelta
import boto3
from botocore.exceptions import ClientError

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('aws_monitor.log'),
        logging.StreamHandler()
    ]
)

class AWSCloudWatch:
    def __init__(self, config_file='config.ini'):
        """初始化AWS CloudWatch客户端"""
        self.config = self._load_config(config_file)
        self.cloudwatch = self._init_client()
        
    def _load_config(self, config_file):
        """加载配置文件"""
        if not os.path.exists(config_file):
            logging.error(f"配置文件不存在: {config_file}")
            sys.exit(1)
            
        config = configparser.ConfigParser()
        config.read(config_file)
        
        if 'aws' not in config:
            logging.error("配置文件中缺少 [aws] 部分")
            sys.exit(1)
            
        return config
        
    def _init_client(self):
        """初始化AWS客户端"""
        try:
            return boto3.client('cloudwatch',
                aws_access_key_id=self.config['aws']['access_key'],
                aws_secret_access_key=self.config['aws']['secret_key'],
                region_name=self.config['aws']['region']
            )
        except Exception as e:
            logging.error(f"初始化AWS客户端失败: {e}")
            sys.exit(1)
            
    def get_metric_statistics(self, namespace, metric_name, dimensions,
                            start_time=None, end_time=None, period=300):
        """获取指标统计数据
        
        Args:
            namespace: 指标命名空间，如 'AWS/EC2'
            metric_name: 指标名称，如 'CPUUtilization'
            dimensions: 维度列表，如 [{'Name': 'InstanceId', 'Value': 'i-1234'}]
            start_time: 开始时间，默认1小时前
            end_time: 结束时间，默认当前时间
            period: 统计周期（秒），默认300秒
        """
        try:
            end_time = end_time or datetime.utcnow()
            start_time = start_time or (end_time - timedelta(hours=1))
            
            response = self.cloudwatch.get_metric_statistics(
                Namespace=namespace,
                MetricName=metric_name,
                Dimensions=dimensions,
                StartTime=start_time,
                EndTime=end_time,
                Period=period,
                Statistics=['Average', 'Maximum']
            )
            return response
            
        except ClientError as e:
            logging.error(f"获取指标统计数据失败: {e}")
            return None
            
    def put_metric_alarm(self, alarm_name, namespace, metric_name,
                        dimensions, threshold, comparison_operator='GreaterThanThreshold',
                        period=300, evaluation_periods=3, alarm_actions=None):
        """创建或更新告警
        
        Args:
            alarm_name: 告警名称
            namespace: 指标命名空间
            metric_name: 指标名称
            dimensions: 维度列表
            threshold: 阈值
            comparison_operator: 比较运算符
            period: 统计周期（秒）
            evaluation_periods: 评估周期数
            alarm_actions: 告警动作列表（SNS主题ARN）
        """
        try:
            self.cloudwatch.put_metric_alarm(
                AlarmName=alarm_name,
                AlarmDescription=f"Auto created alarm for {metric_name}",
                ActionsEnabled=True,
                MetricName=metric_name,
                Namespace=namespace,
                Dimensions=dimensions,
                Period=period,
                EvaluationPeriods=evaluation_periods,
                Threshold=threshold,
                ComparisonOperator=comparison_operator,
                Statistic='Average',
                AlarmActions=alarm_actions or []
            )
            logging.info(f"成功创建/更新告警: {alarm_name}")
            return True
            
        except ClientError as e:
            logging.error(f"创建/更新告警失败: {e}")
            return False
            
    def describe_alarms(self, alarm_names=None):
        """获取告警详情"""
        try:
            if alarm_names:
                response = self.cloudwatch.describe_alarms(
                    AlarmNames=alarm_names
                )
            else:
                response = self.cloudwatch.describe_alarms()
            return response
            
        except ClientError as e:
            logging.error(f"获取告警详情失败: {e}")
            return None
            
    def setup_ec2_monitoring(self):
        """设置EC2实例监控"""
        ec2_config = self.config['ec2_monitor']
        instance_ids = [id.strip() for id in ec2_config['instance_ids'].split(',')]
        
        for instance_id in instance_ids:
            dimensions = [{'Name': 'InstanceId', 'Value': instance_id}]
            
            # CPU使用率告警
            self.put_metric_alarm(
                alarm_name=f"{instance_id}-CPU使用率告警",
                namespace='AWS/EC2',
                metric_name='CPUUtilization',
                dimensions=dimensions,
                threshold=float(ec2_config['cpu_threshold']),
                period=int(ec2_config['check_period']),
                evaluation_periods=int(ec2_config['evaluation_periods']),
                alarm_actions=self.config['alarm_notification']['topic_arns'].split(',')
            )
            
            # 内存使用率告警（需要CloudWatch Agent）
            self.put_metric_alarm(
                alarm_name=f"{instance_id}-内存使用率告警",
                namespace='CWAgent',
                metric_name='mem_used_percent',
                dimensions=dimensions,
                threshold=float(ec2_config['memory_threshold']),
                period=int(ec2_config['check_period']),
                evaluation_periods=int(ec2_config['evaluation_periods']),
                alarm_actions=self.config['alarm_notification']['topic_arns'].split(',')
            )
            
    def setup_rds_monitoring(self):
        """设置RDS实例监控"""
        rds_config = self.config['rds_monitor']
        instance_ids = [id.strip() for id in rds_config['instance_ids'].split(',')]
        
        for instance_id in instance_ids:
            dimensions = [{'Name': 'DBInstanceIdentifier', 'Value': instance_id}]
            
            # CPU使用率告警
            self.put_metric_alarm(
                alarm_name=f"{instance_id}-CPU使用率告警",
                namespace='AWS/RDS',
                metric_name='CPUUtilization',
                dimensions=dimensions,
                threshold=float(rds_config['cpu_threshold']),
                period=300,
                evaluation_periods=3,
                alarm_actions=self.config['alarm_notification']['topic_arns'].split(',')
            )
            
            # 可用存储空间告警
            self.put_metric_alarm(
                alarm_name=f"{instance_id}-存储空间告警",
                namespace='AWS/RDS',
                metric_name='FreeStorageSpace',
                dimensions=dimensions,
                threshold=float(rds_config['disk_threshold']),
                comparison_operator='LessThanThreshold',
                period=300,
                evaluation_periods=3,
                alarm_actions=self.config['alarm_notification']['topic_arns'].split(',')
            )

def main():
    """主函数"""
    try:
        monitor = AWSCloudWatch()
        
        # 设置EC2监控
        monitor.setup_ec2_monitoring()
        
        # 设置RDS监控
        monitor.setup_rds_monitoring()
        
        # 获取所有告警状态
        alarms = monitor.describe_alarms()
        logging.info("当前告警状态:")
        for alarm in alarms.get('MetricAlarms', []):
            logging.info(f"告警名称: {alarm['AlarmName']}, 状态: {alarm['StateValue']}")
            
    except Exception as e:
        logging.error(f"发生错误: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()
